using LinearAlgebra
function LegendreGaussLobatto(p)
    """
    Compute the p'th order LGL solution points at reference domain: [-1,1]
    """
    x = zeros(Float64, p+1)
    
    # p=1, x = [-1, 1]
    if p == 1
        x[1] = -1.0
        x[2] =  1.0
        return x
    end
    
    # p = 2, x = [-1, 0, 1]
    if p == 2
        x[1]= -1.0
        x[2]=  0.0
        x[3]=  1.0   
        return x
    end

    h = @. 2*(0:p-2) + 2
    J = zeros(p-1, p-1)
    @inbounds A = @. sqrt((1:p-2)*((1:p-2)+2)*((1:p-2)+1).*((1:p-2)+1)/(h[1:p-2]+1)/(h[1:p-2]+3))
    @inbounds J[1:p-2,2:p-1] = Diagonal(2.0 ./ (h[1:p-2] .+ 2).*A)
    @inbounds J[1,1] = 0.0
    @inbounds J = J + J'
    
    @inbounds D, V = eigen(J)
    @inbounds i = sortperm(D)
    x = [-1.0; D[i]; 1.0]
end

function Legendre(p, x)
    """
    Evaluate p'th order orhonormal Legendre Polynomial at x
    """
    # Initial values for P(x) at x
    P = zeros(p+1, length(x))
    @. P[1,:] = sqrt(1.0/2.0)
    
    # p = 1
    if  p == 0
        return P'
    end
    
    # p = 2
    @inbounds P[2,:] = sqrt(3.0/2.0)*x
    if  p == 1
        return P[p+1,:]
    end

    # Forward recurrence using the symmetry of the recurrence for p>=2
    a₀ = sqrt(1.0/3.0)
    @simd for i=1:p-1
        a₁ = 2/(2*i+2)*sqrt((i+1)*(i+1)*(i+1)*(i+1)/(2*i+1)/(2*i+3))
        @inbounds P[i+2,:] = 1/a₁*(-a₀*P[i,:] + x.*P[i+1,:])
        a₀ = a₁
    end
    
    return P[p+1,:]
end

function Vandermonde(p, x)
    """
    Initialize the 1D Vandermonde Matrix, V_{ij} = phi_j(r_i)
    """
    V = zeros(length(x), p+1)
    @simd for i in 1:p+1
            V[:,i] = Legendre(i-1, x)
    end
    return V
end

function GradLegendre(p, x)
    """
    Evaluate the derivative of the p'th order Legendre polynomial at x
    """
    P₀ = zeros(length(x),1)
    if p > 0
        @inbounds P₁ = -p*x.*Legendre(p, x) + p*sqrt((2*p+1)/(2*p-1))*Legendre(p-1,x)
        @inbounds P₃ = x.^(p+1)*p*(p+1)/2*sqrt((2*p+1)/2)
        @inbounds I  = abs.(abs.(x) .- 1) .> 10*2.2204e-16
        @inbounds x₁ = x .* I
        @inbounds P₀ = broadcast(~,I).*P₃ + I.*P₁./(1.0 .- x₁.^2)
    end
    return P₀
end

function DiffMatrix(p, x, V)
    """
    Initialize the differentiation matrices D
    """
    Vx = zeros(length(x),p+1)
    for i=0:p
       Vx[:,i+1] = GradLegendre(i, x)
    end

    D = Vx/V
    return D
end